import cv2
import numpy as np
import onnxruntime
import os
from utils.utils import image_transform

model_path = "../resources/model.onnx"
face_feature_path = "../resources/face_feature_vector.txt"
image_paths = "../resources/images_path.txt"
directory_name = "../resources/face_datas"
save_images_path = "../resources/images_path.txt"
end_path = '../resources/face_feature_vector.txt'


def onnx_runtime(img):
    ort_session = onnxruntime.InferenceSession(model_path)
    ort_input1 = {'input': img}
    ort_output1 = ort_session.run(['output'], ort_input1)[0]

    return ort_output1


def face_image_matching(img):
    face_features = np.loadtxt(face_feature_path)
    images_path = []

    with open(image_paths, 'r') as f:
        for line in f.readlines():
            images_path.append(line.split('\n')[0])
    f.close()

    features_distance = []
    for i in range(face_features.shape[0]):
        vector = face_features[i, :]
        vector.reshape(1, 128)
        features_distance.append(np.linalg.norm(img - vector, axis=1)[0])

    return features_distance, images_path


def read_directory(directory_name):
    array_of_img = []
    images_path = []
    for filename in os.listdir(r"./" + directory_name):
        img = cv2.imread(directory_name + "/" + filename)
        images_path.append(directory_name + "/" + filename)
        img = image_transform(img)
        img_dict = {'input': img}
        array_of_img.append(img_dict)
    return array_of_img, images_path


def create_txt(imgs_list):
    face_feature_vector = []
    ort_session = onnxruntime.InferenceSession(model_path)

    for img_dict in imgs_list:
        ort_output = ort_session.run(['output'], img_dict)[0]
        face_feature_vector.append(ort_output.ravel())

    face_feature_vector = np.array(face_feature_vector)
    # 将人脸图片库的facenet运算后的128维特征向量保存到face_feature_vector.txt
    np.savetxt(end_path, face_feature_vector, )


def write_images_path(path):
    f = open(save_images_path, 'w')
    for line in path:
        f.write(line + '\n')
    f.close()


def update_dates():
    array_of_img, images_path = read_directory(directory_name)
    write_images_path(images_path)
    create_txt(array_of_img)
